//******************************************************************************
//
// Copyright (c) Microsoft. All rights reserved.
//
// This code is licensed under the MIT License (MIT).
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
//
//******************************************************************************

#import "NSKeyValueObserving-Internal.h"
#import "NSObject_NSKeyValueCoding-Internal.h"

#import "NSMethodSignatureInternal.h"
#import <objc/encoding.h>
#import <objc/runtime.h>
#import "type_encoding_cases.h"

#import <memory>

/* These are defined by the ABI and the runtime. */
#define ABI_SUPER(obj) (((Class**)obj)[0][1])
#define ABI_ISA(obj) (((Class*)obj)[0])

#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wdeprecated-objc-pointer-introspection"
// Small objects, as defined by libobjc2, are tagged pointers. Instead of being heap-allocated,
// the object is stored in the sizeof(id) - OBJC_SMALL_OBJECT_BITS bits of the "pointer" itself.
static inline bool isSmallObject_np(id object) {
    return ((((uintptr_t)object) & OBJC_SMALL_OBJECT_MASK) != 0);
}
#pragma clang diagnostic pop

static const wchar_t* TAG = L"NSObject (NSKeyValueObservation)";

static void NSKVO$setObject$forKey$(id self, SEL _cmd, id object, NSString* key);
static void NSKVO$removeObjectForKey$(id self, SEL _cmd, NSString* key);
static void NSKVO$nilIMP(id self, SEL _cmd) {
}

static void _NSKVOEnsureObjectIsKVOAware(id object) {
    // We have to use ABI_ISA here, because object_getClass will skip the secret hidden subclass.
    if (class_respondsToSelector(ABI_ISA(object), @selector(_isKVOAware))) {
        // The presence of _isKVOAware signals that we have already mangled this object.
        return;
    }

    /*
     A crucial design decision was made here: because object_getClass() skips autogenerated subclasses, the likes of which are used for
     associated objects, the KVO machinery (if implemented using manual subclassing) would delete all associated objects on an observed
     instance. That was deemed unacceptable.

     Likewise, this implementation is not free of issues:
     - The _np methods are nonportable.
     - Anyone using class_getMethodImplementation(object_getClass(...), ...) will receive the original IMP for any overridden method.
     - We have to manually load isa to get at the secret subclass (thus the use of ABI_ISA/ABI_SUPER.)
     - It is dependent upon a libobjc2 implementation detail:
         object_addMethod_np creates a hidden subclass for the object's class (one per object!)
    */

    object_addMethod_np(object, @selector(setObject:forKey:), reinterpret_cast<IMP>(NSKVO$setObject$forKey$), "v@:@@");
    object_addMethod_np(object, @selector(removeObjectForKey:), reinterpret_cast<IMP>(NSKVO$removeObjectForKey$), "v@:@");

    object_addMethod_np(object, @selector(_isKVOAware), reinterpret_cast<IMP>(NSKVO$nilIMP), "v@:");
}

#pragma region Method Implementations
// Selector mappings: the class-level mapping from a selector (setX:) to the KVC key ("x")
// to which it corresponds. This is necessary because both "X" and "x" map to setX:, but we need
// to be cognizant of precisely which it was for any given observee.
// The sole reason we can get away with keeping selector<->key mappings on the class is that,
// through the lifetime of said class, it can never lose selectors. This mapping will be pertinent
// to every instance of the class.
static inline CFMutableDictionaryRef _selectorMappingsForObject(id object) {
    static char s_selMapKey;
    auto cls = object_getClass(object); // here we explicitly want the public (non-hidden) class associated with the object.
    @synchronized(cls) {
        CFMutableDictionaryRef selMappings = (CFMutableDictionaryRef)objc_getAssociatedObject(cls, &s_selMapKey);
        if (!selMappings) {
            selMappings = CFDictionaryCreateMutable(kCFAllocatorDefault,
                                                    16,
                                                    NULL,
                                                    &kCFTypeDictionaryValueCallBacks); // no key callbacks: they're selector names.
            objc_setAssociatedObject(cls, &s_selMapKey, (id)selMappings, OBJC_ASSOCIATION_RETAIN);
            CFRelease(selMappings);
        }
        return selMappings;
    }
}

static inline NSString* _keyForSelector(id object, SEL selector) {
    return (NSString*)CFDictionaryGetValue(_selectorMappingsForObject(object), sel_getName(selector));
}

static inline void _associateSelectorWithKey(id object, SEL selector, NSString* key) {
    CFDictionarySetValue(_selectorMappingsForObject(object), sel_getName(selector), key);
}

template <typename T>
static void notifyingSetImpl(id self, SEL _cmd, T val) {
    NSString* key = _keyForSelector(self, _cmd);

    [self willChangeValueForKey:key];

    // This matches what clang generates for [super xxx].
    struct objc_super super {
        self, ABI_SUPER(self)
    };
    auto imp = (void (*)(id, SEL, T))objc_msg_lookup_super(&super, _cmd);
    imp(self, _cmd, val);

    [self didChangeValueForKey:key];
}

static void notifyingVariadicSetImpl(id self, SEL _cmd, ...) {
    NSString* key = _keyForSelector(self, _cmd);

    // [Source: NSInvocation.mm]
    // This attempts to flatten the method's arguments (as determined by its type encoding)
    // from the stack into a buffer. That buffer is then emitted back onto the stack for the imp callthrough.
    // This only works if we assume that our calling convention passes variadics and non-variadics
    // in the same way: on the stack. For our two supported platforms, this seems to hold true.
    NSMethodSignature* sig = [self methodSignatureForSelector:_cmd];
    auto argSz = [sig getArgumentSizeAtIndex:2];
    auto nStackArgs = argSz / sizeof(uintptr_t);
    uintptr_t* raw = static_cast<uintptr_t*>(calloc(sizeof(uintptr_t), nStackArgs));
    va_list ap;
    va_start(ap, _cmd);
    for (uintptr_t i = 0; i < nStackArgs; ++i) {
        raw[i] = va_arg(ap, uintptr_t);
    }
    va_end(ap);

    [self willChangeValueForKey:key];

    struct objc_super super {
        self, ABI_SUPER(self)
    };
    auto imp = (void (*)(id, SEL, ...))objc_msg_lookup_super(&super, _cmd);

    // VSO 5955259; NSInvocation informs this implementation and this will need
    // to be cleaned up when NSInvocation is.
    switch (nStackArgs) {
        case 1:
            imp(self, _cmd, raw[0]);
            break;
        case 2:
            imp(self, _cmd, raw[0], raw[1]);
            break;
        case 3:
            imp(self, _cmd, raw[0], raw[1], raw[2]);
            break;
        case 4:
            imp(self, _cmd, raw[0], raw[1], raw[2], raw[3]);
            break;
        case 5:
            imp(self, _cmd, raw[0], raw[1], raw[2], raw[3], raw[4]);
            break;
        case 6:
            imp(self, _cmd, raw[0], raw[1], raw[2], raw[3], raw[4], raw[5]);
            break;
        default:
            FAIL_FAST_MSG("Can't override setter with more than 6 sizeof(long int) stack arguments.");
            return;
    }

    [self didChangeValueForKey:key];

    free(raw);
}

static void NSKVONotifying$insertObject$inXxxAtIndex$(id self, SEL _cmd, id object, NSUInteger index) {
    NSString* key = _keyForSelector(self, _cmd);
    NSIndexSet* indexes = [NSIndexSet indexSetWithIndex:index];

    [self willChange:NSKeyValueChangeInsertion valuesAtIndexes:indexes forKey:key];

    objc_super super{ self, ABI_SUPER(self) };
    auto imp = (void (*)(id, SEL, id, NSUInteger))objc_msg_lookup_super(&super, _cmd);
    imp(self, _cmd, object, index);

    [self didChange:NSKeyValueChangeInsertion valuesAtIndexes:indexes forKey:key];
}

static void NSKVONotifying$insertXxx$atIndexes$(id self, SEL _cmd, id object, NSIndexSet* indexes) {
    NSString* key = _keyForSelector(self, _cmd);

    [self willChange:NSKeyValueChangeInsertion valuesAtIndexes:indexes forKey:key];

    objc_super super{ self, ABI_SUPER(self) };
    auto imp = (void (*)(id, SEL, id, NSIndexSet*))objc_msg_lookup_super(&super, _cmd);
    imp(self, _cmd, object, indexes);

    [self didChange:NSKeyValueChangeInsertion valuesAtIndexes:indexes forKey:key];
}

static void NSKVONotifying$removeObjectFromXxxAtIndex$(id self, SEL _cmd, NSUInteger index) {
    NSString* key = _keyForSelector(self, _cmd);
    NSIndexSet* indexes = [NSIndexSet indexSetWithIndex:index];

    [self willChange:NSKeyValueChangeRemoval valuesAtIndexes:indexes forKey:key];

    objc_super super{ self, ABI_SUPER(self) };
    auto imp = (void (*)(id, SEL, NSUInteger))objc_msg_lookup_super(&super, _cmd);
    imp(self, _cmd, index);

    [self didChange:NSKeyValueChangeRemoval valuesAtIndexes:indexes forKey:key];
}

static void NSKVONotifying$removeXxxAtIndexes$(id self, SEL _cmd, NSIndexSet* indexes) {
    NSString* key = _keyForSelector(self, _cmd);

    [self willChange:NSKeyValueChangeRemoval valuesAtIndexes:indexes forKey:key];

    objc_super super{ self, ABI_SUPER(self) };
    auto imp = (void (*)(id, SEL, NSIndexSet*))objc_msg_lookup_super(&super, _cmd);
    imp(self, _cmd, indexes);

    [self didChange:NSKeyValueChangeRemoval valuesAtIndexes:indexes forKey:key];
}

static void NSKVONotifying$replaceObjectInXxxAtIndex$withObject$(id self, SEL _cmd, NSUInteger index, id object) {
    NSString* key = _keyForSelector(self, _cmd);
    NSIndexSet* indexes = [NSIndexSet indexSetWithIndex:index];

    [self willChange:NSKeyValueChangeReplacement valuesAtIndexes:indexes forKey:key];

    objc_super super{ self, ABI_SUPER(self) };
    auto imp = (void (*)(id, SEL, NSUInteger, id))objc_msg_lookup_super(&super, _cmd);
    imp(self, _cmd, index, object);

    [self didChange:NSKeyValueChangeReplacement valuesAtIndexes:indexes forKey:key];
}

static void NSKVONotifying$replaceXxxAtIndexes$withXxx$(id self, SEL _cmd, NSIndexSet* indexes, NSArray* objects) {
    NSString* key = _keyForSelector(self, _cmd);

    [self willChange:NSKeyValueChangeReplacement valuesAtIndexes:indexes forKey:key];

    objc_super super{ self, ABI_SUPER(self) };
    auto imp = (void (*)(id, SEL, NSIndexSet*, NSArray*))objc_msg_lookup_super(&super, _cmd);
    imp(self, _cmd, indexes, objects);

    [self didChange:NSKeyValueChangeReplacement valuesAtIndexes:indexes forKey:key];
}

template <NSKeyValueSetMutationKind Kind>
static inline void _NSKVOSetDispatch(id self, SEL _cmd, NSSet* set) {
    NSString* key = _keyForSelector(self, _cmd);

    [self willChangeValueForKey:key withSetMutation:Kind usingObjects:set];

    objc_super super{ self, ABI_SUPER(self) };
    auto imp = (void (*)(id, SEL, NSSet*))objc_msg_lookup_super(&super, _cmd);
    imp(self, _cmd, set);

    [self didChangeValueForKey:key withSetMutation:Kind usingObjects:set];
}

template <NSKeyValueSetMutationKind Kind>
static inline void _NSKVOSetDispatchIndividual(id self, SEL _cmd, id obj) {
    NSSet* set = [NSSet setWithObject:obj];

    NSString* key = _keyForSelector(self, _cmd);
    [self willChangeValueForKey:key withSetMutation:Kind usingObjects:set];

    objc_super super{ self, ABI_SUPER(self) };
    auto imp = (void (*)(id, SEL, id))objc_msg_lookup_super(&super, _cmd);
    imp(self, _cmd, obj);

    [self didChangeValueForKey:key withSetMutation:Kind usingObjects:set];
}

// - (void)setObject:(id)object forKey:(NSString*)key
static void NSKVO$setObject$forKey$(id self, SEL _cmd, id object, NSString* key) {
    [self willChangeValueForKey:key];
    objc_super super{ self, ABI_SUPER(self) };
    auto imp = (void (*)(id, SEL, id, id))objc_msg_lookup_super(&super, _cmd);
    imp(self, _cmd, object, key);
    [self didChangeValueForKey:key];
}

// - (void)removeObjectForKey:(NSString*)key
static void NSKVO$removeObjectForKey$(id self, SEL _cmd, NSString* key) {
    [self willChangeValueForKey:key];
    objc_super super{ self, ABI_SUPER(self) };
    auto imp = (void (*)(id, SEL, id))objc_msg_lookup_super(&super, _cmd);
    imp(self, _cmd, key);
    [self didChangeValueForKey:key];
}
#pragma endregion

#define KVO_SET_IMPL_CASE(type, name, capitalizedName, encodingChar) \
    case encodingChar:                                               \
        newImpl = reinterpret_cast<IMP>(&notifyingSetImpl<type>);    \
        break;

// invariant: rawKey has already been capitalized
static inline void _NSKVOEnsureSimpleKeyWillNotify(id object, NSString* key, const char* rawKey) {
    auto sel = KVCSetterForPropertyName(object, rawKey);

    Method originalMethod = class_getInstanceMethod(object_getClass(object), sel);
    const char* types = method_getTypeEncoding(originalMethod);
    if (!types) {
        return;
    }

    NSMethodSignature* sig = [NSMethodSignature signatureWithObjCTypes:types];

    auto valueType = [sig getArgumentTypeAtIndex:2];
    IMP newImpl = NULL;

    switch (valueType[0]) {
        OBJC_APPLY_ALL_TYPE_ENCODINGS(KVO_SET_IMPL_CASE);
        case '{':
        case '[': {
            size_t valueSize = objc_sizeof_type(valueType);
            if (valueSize > 6 * sizeof(uintptr_t)) {
                [NSException raise:NSInvalidArgumentException
                            format:@"Class %s key %@ has a value size of %u bytes, and cannot currently be KVO compliant.",
                                   class_getName(object_getClass(object)),
                                   key,
                                   static_cast<uintptr_t>(valueSize)];
            }
            newImpl = reinterpret_cast<IMP>(&notifyingVariadicSetImpl);
            break;
        }
        default:
            [NSException raise:NSInvalidArgumentException
                        format:@"Class %s is not KVO compliant for key %@.", class_getName(object_getClass(object)), key];
            return;
    }

    _associateSelectorWithKey(object, sel, key);
    object_addMethod_np(object, sel, newImpl, types);
}

static void replaceAndAssociateWithKey(id object, SEL sel, NSString* key, IMP imp) {
    if ([object respondsToSelector:sel]) {
        const char* selName = sel_getName(sel);
        Method method = class_getInstanceMethod(object_getClass(object), sel);
        if (!method) {
            TraceWarning(TAG,
                         L"Unable to find method for %hs on class %hs; perhaps it is a forward?",
                         selName,
                         object_getClassName(object));
            return;
        }

        _associateSelectorWithKey(object, sel, key);
        object_replaceMethod_np(object, sel, imp, method_getTypeEncoding(method));
    }
}

static SEL formatSelector(NSString* format, ...) {
    va_list ap;
    va_start(ap, format);
    SEL sel = NSSelectorFromString([[[NSString alloc] initWithFormat:format arguments:ap] autorelease]);
    va_end(ap);
    return sel;
}

// invariant: rawKey has already been capitalized
static inline void _NSKVOEnsureOrderedCollectionWillNotify(id object, NSString* key, const char* rawKey) {
    SEL insertOne = formatSelector(@"insertObject:in%sAtIndex:", rawKey);
    SEL insertMany = formatSelector(@"insert%s:atIndexes:", rawKey);
    if ([object respondsToSelector:insertOne] || [object respondsToSelector:insertMany]) {
        replaceAndAssociateWithKey(object, insertOne, key, (IMP)NSKVONotifying$insertObject$inXxxAtIndex$);
        replaceAndAssociateWithKey(object, insertMany, key, (IMP)NSKVONotifying$insertXxx$atIndexes$);
        replaceAndAssociateWithKey(object,
                                   formatSelector(@"removeObjectFrom%sAtIndex:", rawKey),
                                   key,
                                   (IMP)NSKVONotifying$removeObjectFromXxxAtIndex$);
        replaceAndAssociateWithKey(object, formatSelector(@"remove%sAtIndexes:", rawKey), key, (IMP)NSKVONotifying$removeXxxAtIndexes$);
        replaceAndAssociateWithKey(object,
                                   formatSelector(@"replaceObjectIn%sAtIndex:withObject:", rawKey),
                                   key,
                                   (IMP)NSKVONotifying$replaceObjectInXxxAtIndex$withObject$);
        replaceAndAssociateWithKey(object,
                                   formatSelector(@"replace%sAtIndexes:with%s:", rawKey, rawKey),
                                   key,
                                   (IMP)NSKVONotifying$replaceXxxAtIndexes$withXxx$);
    }
}

// invariant: rawKey has already been capitalized
static inline void _NSKVOEnsureUnorderedCollectionWillNotify(id object, NSString* key, const char* rawKey) {
    SEL addOne = formatSelector(@"add%sObject:", rawKey);
    SEL addMany = formatSelector(@"add%s:", rawKey);
    SEL removeOne = formatSelector(@"remove%sObject:", rawKey);
    SEL removeMany = formatSelector(@"remove%s:", rawKey);
    if (([object respondsToSelector:addOne] || [object respondsToSelector:addMany]) &&
        ([object respondsToSelector:removeOne] || [object respondsToSelector:removeMany])) {
        replaceAndAssociateWithKey(object, addOne, key, (IMP)_NSKVOSetDispatchIndividual<NSKeyValueUnionSetMutation>);
        replaceAndAssociateWithKey(object, addMany, key, (IMP)_NSKVOSetDispatch<NSKeyValueUnionSetMutation>);
        replaceAndAssociateWithKey(object, removeOne, key, (IMP)_NSKVOSetDispatchIndividual<NSKeyValueMinusSetMutation>);
        replaceAndAssociateWithKey(object, removeMany, key, (IMP)_NSKVOSetDispatch<NSKeyValueMinusSetMutation>);
        replaceAndAssociateWithKey(object,
                                   formatSelector(@"intersect%s:", rawKey),
                                   key,
                                   (IMP)_NSKVOSetDispatch<NSKeyValueIntersectSetMutation>);
    }
}

static std::unique_ptr<char[]> mutableBufferFromString(NSString* string) {
    NSUInteger lengthInBytes = 0;
    auto lengthInCharacters = string.length;
    [string getBytes:nullptr
             maxLength:NSIntegerMax
            usedLength:&lengthInBytes
              encoding:NSUTF8StringEncoding
               options:0
                 range:(NSRange) { 0, lengthInCharacters }
        remainingRange:nullptr];
    std::unique_ptr<char[]> rawKey = std::make_unique<char[]>(lengthInBytes + 1);
    [string getBytes:rawKey.get()
             maxLength:lengthInBytes
            usedLength:&lengthInBytes
              encoding:NSUTF8StringEncoding
               options:0
                 range:(NSRange) { 0, lengthInCharacters }
        remainingRange:nullptr];
    rawKey[lengthInBytes] = '\0';
    return rawKey;
}

// NSKVOEnsureKeyWillNotify is the main entrypoint into the swizzling code.
void _NSKVOEnsureKeyWillNotify(id object, NSString* key) {
    // Since we cannot replace the isa of tagged pointer objects, we can't swizzle them.
    if (isSmallObject_np(object)) {
        return;
    }

    // A class is allowed to decline automatic swizzling for any/all of its keys.
    if (![[object class] automaticallyNotifiesObserversForKey:key]) {
        return;
    }

    std::unique_ptr<char[]> rawKey{ std::move(mutableBufferFromString(key)) };
    rawKey[0] = toupper(rawKey[0]);

    @synchronized(object) {
        _NSKVOEnsureObjectIsKVOAware(object);
        _NSKVOEnsureSimpleKeyWillNotify(object, key, rawKey.get());
        _NSKVOEnsureOrderedCollectionWillNotify(object, key, rawKey.get());
        _NSKVOEnsureUnorderedCollectionWillNotify(object, key, rawKey.get());
    }
}
